/*
 * Licensed to the Apache Software Foundation (ASF) under one or more contributor license agreements. See the NOTICE file distributed with this work for additional information regarding copyright
 * ownership. The ASF licenses this file to You under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance with the License. You may obtain a copy of the
 * License at
 *
 * http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software distributed under the License is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
 * implied. See the License for the specific language governing permissions and limitations under the License.
 *
 */

package com.spark.sniffer.plugin.interceptor.enhance;

import com.spark.sniffer.plugin.AbstractClassEnhancePluginDefine;
import com.spark.sniffer.plugin.interceptor.ConstructorInterceptPoint;
import com.spark.sniffer.plugin.interceptor.DeclaredInstanceMethodsInterceptPoint;
import com.spark.sniffer.plugin.interceptor.InstanceMethodsInterceptPoint;
import com.spark.sniffer.plugin.interceptor.StaticMethodsInterceptPoint;
import net.bytebuddy.asm.Advice;
import net.bytebuddy.description.method.MethodDescription;
import net.bytebuddy.description.type.TypeDescription;
import net.bytebuddy.dynamic.DynamicType;
import net.bytebuddy.implementation.MethodDelegation;
import net.bytebuddy.implementation.SuperMethodCall;
import net.bytebuddy.implementation.bind.annotation.Morph;
import net.bytebuddy.matcher.ElementMatcher;
import net.bytebuddy.matcher.ElementMatchers;

import static net.bytebuddy.matcher.ElementMatchers.*;

/**
 * This class controls all enhance operations, including enhance constructors, instance methods and static methods. All the enhances base on three types interceptor point:
 * {@link ConstructorInterceptPoint}, {@link InstanceMethodsInterceptPoint} and {@link StaticMethodsInterceptPoint} If plugin is going to enhance constructors, instance methods, or both,
 * {@link ClassEnhancePluginDefine} will add a field of {@link Object} type.
 *
 * @author wusheng
 */
public abstract class ClassEnhancePluginDefine extends AbstractClassEnhancePluginDefine {
    /**
     * New field name.
     */
    public static final String CONTEXT_ATTR_NAME = "_$EnhancedClassField_ws";

    /**
     * Begin to define how to enhance class. After invoke this method, only means definition is finished.
     *
     * @param typeDescription target class description
     * @param newClassBuilder byte-buddy's builder to manipulate class bytecode.
     * @return new byte-buddy's builder for further manipulation.
     */
    @Override
    protected DynamicType.Builder<?> enhance(TypeDescription typeDescription, DynamicType.Builder<?> newClassBuilder,
        ClassLoader classLoader) throws Exception {
        newClassBuilder = this.enhanceClass(typeDescription, newClassBuilder, classLoader);
        newClassBuilder = this.enhanceInstance(typeDescription, newClassBuilder, classLoader);

        return newClassBuilder;
    }

    /**
     * Enhance a class to intercept class static methods.
     *
     * @param typeDescription target class description
     * @param newClassBuilder byte-buddy's builder to manipulate class bytecode.
     * @return new byte-buddy's builder for further manipulation.
     */
    private DynamicType.Builder<?> enhanceClass(TypeDescription typeDescription, DynamicType.Builder<?> newClassBuilder,
        ClassLoader classLoader) throws Exception {
        StaticMethodsInterceptPoint[] staticMethodsInterceptPoints = getStaticMethodsInterceptPoints();
        String enhanceOriginClassName = typeDescription.getTypeName();
        if (staticMethodsInterceptPoints == null || staticMethodsInterceptPoints.length == 0) {
            return newClassBuilder;
        }
        if(isBootstrapInstrument()) {
            for (StaticMethodsInterceptPoint staticMethodsInterceptPoint : staticMethodsInterceptPoints) {
                String interceptor = staticMethodsInterceptPoint.getMethodsInterceptor().getName().replace("/", ".");
                if (interceptor == null) {
                    throw new Exception(
                        "no InstanceMethodsAroundInterceptor define to enhance class " + enhanceOriginClassName);
                }
                ElementMatcher.Junction<MethodDescription> junction =
                    isStatic().and(staticMethodsInterceptPoint.getMethodsMatcher());
                newClassBuilder = newClassBuilder.visit(Advice.to(staticMethodsInterceptPoint.getMethodsInterceptor())
                    .on(junction));
            }
        } else {
            for (StaticMethodsInterceptPoint staticMethodsInterceptPoint : staticMethodsInterceptPoints) {
                String interceptor = staticMethodsInterceptPoint.getMethodsInterceptor().getName().replace("/", ".");
                if (interceptor == null || interceptor.trim().length() == 0) {
                    throw new Exception(
                        "no StaticMethodsAroundInterceptor define to enhance class " + enhanceOriginClassName);
                }

                if (staticMethodsInterceptPoint.isOverrideArgs()) {
                    newClassBuilder =
                        newClassBuilder.method(isStatic().and(staticMethodsInterceptPoint.getMethodsMatcher()))
                            .intercept(MethodDelegation.withDefaultConfiguration()
                                .withBinders(Morph.Binder.install(OverrideCallable.class))
                                .to(new StaticMethodsInterWithOverrideArgs(interceptor)));
                } else {
                    newClassBuilder =
                        newClassBuilder.method(isStatic().and(staticMethodsInterceptPoint.getMethodsMatcher()))
                            .intercept(MethodDelegation.withDefaultConfiguration().to(new StaticMethodsInter(interceptor)));
                }

            }
        }
        return newClassBuilder;
    }

    /**
     * Enhance a class to intercept constructors and class instance methods.
     *
     * @param typeDescription target class description
     * @param newClassBuilder byte-buddy's builder to manipulate class bytecode.
     * @return new byte-buddy's builder for further manipulation.
     */
    private DynamicType.Builder<?> enhanceInstance(TypeDescription typeDescription,
        DynamicType.Builder<?> newClassBuilder, ClassLoader classLoader) throws Exception {
        InstanceMethodsInterceptPoint[] instanceMethodsInterceptPoints = getInstanceMethodsInterceptPoints();
        ConstructorInterceptPoint[] constructorInterceptPoints = getConstructorsInterceptPoints();
        String enhanceOriginClassName = typeDescription.getTypeName();
        boolean existedMethodsInterceptPoints = false;
        boolean existedConstructorInterceptPoint = false;
        if (constructorInterceptPoints != null && constructorInterceptPoints.length > 0) {
            existedConstructorInterceptPoint = true;
        }
        if (instanceMethodsInterceptPoints != null && instanceMethodsInterceptPoints.length > 0) {
            existedMethodsInterceptPoints = true;
        }
        /**
         * nothing need to be enhanced in class instance, maybe need enhance static methods.
         */
        if (!existedMethodsInterceptPoints && !existedConstructorInterceptPoint) {
            return newClassBuilder;
        }

        /**
         * 2. enhance constructors
         */
        if (existedConstructorInterceptPoint) {
            if(isBootstrapInstrument()) {
                for (ConstructorInterceptPoint constructorInterceptPoint : constructorInterceptPoints) {
                    String interceptor = constructorInterceptPoint.getConstructorInterceptor().getName().replace("/", ".");
                    if (interceptor == null) {
                        throw new Exception(
                            "no InstanceMethodsAroundInterceptor define to enhance class " + enhanceOriginClassName);
                    }
                    ElementMatcher.Junction<MethodDescription> junction = isConstructor().and(constructorInterceptPoint.getConstructorMatcher());
                    newClassBuilder = newClassBuilder.visit(Advice.to(constructorInterceptPoint.getConstructorInterceptor())
                        .on(junction));
                }
            } else {
                for (ConstructorInterceptPoint constructorInterceptPoint : constructorInterceptPoints) {
                    String interceptor = constructorInterceptPoint.getConstructorInterceptor().getName().replace("/", ".");
                    newClassBuilder = newClassBuilder.constructor(constructorInterceptPoint.getConstructorMatcher())
                        .intercept(SuperMethodCall.INSTANCE.andThen(MethodDelegation.withDefaultConfiguration()
                            .to(new ConstructorInter(interceptor, classLoader))));
                }
            }
        }

        /**
         * 3. enhance instance methods
         */
        if(existedMethodsInterceptPoints) {
            if(isBootstrapInstrument()) {
                for (InstanceMethodsInterceptPoint instanceMethodsInterceptPoint : instanceMethodsInterceptPoints) {
                    String interceptor = instanceMethodsInterceptPoint.getMethodsInterceptor().getName().replace("/", ".");
                    if (interceptor == null) {
                        throw new Exception(
                            "no InstanceMethodsAroundInterceptor define to enhance class " + enhanceOriginClassName);
                    }
                    ElementMatcher.Junction<MethodDescription> junction =
                        not(isStatic()).and(instanceMethodsInterceptPoint.getMethodsMatcher());
                    if (instanceMethodsInterceptPoint instanceof DeclaredInstanceMethodsInterceptPoint) {
                        junction = junction.and(ElementMatchers.<MethodDescription>isDeclaredBy(typeDescription));
                    }
                    newClassBuilder = newClassBuilder.visit(Advice.to(instanceMethodsInterceptPoint.getMethodsInterceptor())
                        .on(junction));
                }
            } else {
                for (InstanceMethodsInterceptPoint instanceMethodsInterceptPoint : instanceMethodsInterceptPoints) {
                    String interceptor = instanceMethodsInterceptPoint.getMethodsInterceptor().getName().replace("/", ".");
                    if (interceptor == null) {
                        throw new Exception(
                            "no InstanceMethodsAroundInterceptor define to enhance class " + enhanceOriginClassName);
                    }
                    ElementMatcher.Junction<MethodDescription> junction =
                        not(isStatic()).and(instanceMethodsInterceptPoint.getMethodsMatcher());
                    if (instanceMethodsInterceptPoint instanceof DeclaredInstanceMethodsInterceptPoint) {
                        junction = junction.and(ElementMatchers.<MethodDescription>isDeclaredBy(typeDescription));
                    }
                    if (instanceMethodsInterceptPoint.isOverrideArgs()) {
                        newClassBuilder = newClassBuilder.method(junction)
                            .intercept(MethodDelegation.withDefaultConfiguration()
                                .withBinders(Morph.Binder.install(OverrideCallable.class))
                                .to(new InstMethodsInterWithOverrideArgs(interceptor, classLoader)));
                    } else {
                        newClassBuilder = newClassBuilder.method(junction).intercept(
                            MethodDelegation.withDefaultConfiguration().to(new InstMethodsInter(interceptor, classLoader)));
                    }
                }
            }
        }
        return newClassBuilder;
    }

    /**
     * Instance methods intercept point. See {@link InstanceMethodsInterceptPoint}
     *
     * @return collections of {@link InstanceMethodsInterceptPoint}
     */
    protected abstract InstanceMethodsInterceptPoint[] getInstanceMethodsInterceptPoints();

    protected abstract StaticMethodsInterceptPoint[] getStaticMethodsInterceptPoints();

    /**
     * Constructor methods intercept point. See {@link ConstructorInterceptPoint}
     *
     * @return collections of {@link ConstructorInterceptPoint}
     */
    public abstract ConstructorInterceptPoint[] getConstructorsInterceptPoints();

}
